import cv2
import numpy as np
from sklearn.svm import SVC
import cv2_pic_trans
import os
import pickle
import random
import joblib


def read_pic(dir_path, iter, batch_size):
    data = []
    label = []
    file_list = os.listdir(dir_path)
    random.shuffle(file_list)
    length = len(file_list)
    '''
    0 ~ 99
    100 ~ 199
    200 ~ 299
    '''
    if iter * batch_size > length - 1:
        return data, label

    for i in range(iter * batch_size, (iter + 1) * batch_size - 1):
        if i > length - 1:
            break
        clazz = file_list[i].split(".")[0]
        if clazz == "cat":
            label.append(0)
        else:
            label.append(1)
        image_path = os.path.join(dir_path, file_list[i])
        res = cv2_pic_trans.image_trans(image_path)
        data.append(res)
    return data, label

def train(train_dir, batch_size):

    clf = SVC()
    iter = 0
    while True:
        train_data, train_labels = read_pic(train_dir, iter, batch_size)
        if len(train_data) == 0:
            break
        clf.fit(np.array(train_data), np.array(train_labels))
        iter += 1
        if iter % 10 == 0:
            print("{}次batch训练完成".format(iter))
        if iter == 1:
            break

    model = pickle.dumps(clf)
    f = open('svm.model', "wb+")
    f.write(model)
    f.close()

def test(test_dir):

    with open('svm.model', 'rb') as f:
        svm = pickle.load(f)

        file_list = os.listdir(test_dir)
        test_data = []
        test_label = [1, 1, 1, 0, 0, 0]
        for file in file_list:
            image_path = os.path.join(test_dir, file)
            res = cv2_pic_trans.image_trans(image_path)
            test_data.append(res)

        print(svm.predict(np.array(test_data)))
        print(test_label)


if __name__ == "__main__":
    train_dir = "D:\\workstation\\test_pytorch\\DogsVSCats\\data\\train"
    test_dir = "D:\\workstation\\test_pytorch\\DogsVSCats\\data\\mini_test"

    batch_size = 10000
    train(train_dir, batch_size)
    test(test_dir)

